import numpy as np
import cv2
import math
from matplotlib import pyplot as plt


def make_transform_complex_matrix(nums, params):
    """
    傅里叶正变换矩阵形式的变换矩阵G
    :param nums: 参数，表示height或width
    :param param: DFT 或 IDFT
    :return: 变换矩阵G 数据类型np.complex64
    """
    row = np.array([np.arange(nums)])
    baseMat = np.dot(row.T, row).astype(np.complex64)
    if params == 'DFT':
        return np.exp((-2j * np.pi / nums) * baseMat)
    elif params == 'IDFT':
        return np.exp((2j * np.pi / nums) * baseMat)


def make_transform_complex_matrix_2(nums, params, lr):
    """
    傅里叶正变换矩阵形式的变换矩阵G
    :param nums: 参数，表示height或width
    :param param: DFT 或 IDFT
    :return: 变换矩阵G 数据类型np.complex64
    """
    l_dft = 100
    l_raw = 200
    if params == 'DFT':
        if lr == 'l':
            row = np.array([np.arange(l_dft)])
            column = np.array([np.arange(nums)])
            baseMat = np.dot(row.T, column).astype(np.complex64)
            return np.exp((-2j * np.pi / l_dft) * baseMat)
        else:
            row = np.array([np.arange(nums)])
            column = np.array([np.arange(l_dft)])
            baseMat = np.dot(row.T, column).astype(np.complex64)
            return np.exp((-2j * np.pi / l_dft) * baseMat)
    elif params == 'IDFT':
        if lr == 'l':
            row = np.array([np.arange(l_raw)])
            column = np.array([np.arange(l_dft)])
            baseMat = np.dot(row.T, column).astype(np.complex64)
            return np.exp((2j * np.pi / l_dft) * baseMat)
        else:
            row = np.array([np.arange(l_dft)])
            column = np.array([np.arange(l_raw)])
            baseMat = np.dot(row.T, column).astype(np.complex64)
            return np.exp((-2j * np.pi / l_dft) * baseMat)


def dft_complex_2(img):
    """
    二维离散傅里叶变换，直接使用复矩阵相乘实现
    :param img: 一张灰度图 H*W
    :return: img_dft H*W np.complex64
    """
    img = np.float32(img)
    height, width = img.shape
    Gl = make_transform_complex_matrix(height, params='DFT', lr='l')
    Gr = make_transform_complex_matrix(width, params='DFT', lr='r')
    img_dft = np.dot(np.dot(Gl, img), Gr)
    return img_dft


def idft_complex_2(img_dft):
    """
    二维离散傅里叶逆变换，直接使用复矩阵相乘实现
    :param img_dft: DFT的变换结果，H*W，np.complex64
    :return: IDFT的变换结果，H*W，np.complex64
    """
    height, width = img_dft.shape
    Gl = make_transform_complex_matrix(height, params='IDFT', lr='l')
    Gr = make_transform_complex_matrix(width, params='IDFT', lr='r')
    img_idft = np.dot(np.dot(Gl, img_dft), Gr) / (height * width)
    return img_idft


def dft_complex(img):
    """
    二维离散傅里叶变换，直接使用复矩阵相乘实现
    :param img: 一张灰度图 H*W
    :return: img_dft H*W np.complex64
    """
    img = np.float32(img)
    height, width = img.shape
    Gl = make_transform_complex_matrix(height, params='DFT')
    Gr = make_transform_complex_matrix(width, params='DFT')
    img_dft = np.dot(np.dot(Gl, img), Gr)
    return img_dft


def idft_complex(img_dft):
    """
    二维离散傅里叶逆变换，直接使用复矩阵相乘实现
    :param img_dft: DFT的变换结果，H*W，np.complex64
    :return: IDFT的变换结果，H*W，np.complex64
    """
    height, width = img_dft.shape
    Gl = make_transform_complex_matrix(height, params='IDFT')
    Gr = make_transform_complex_matrix(width, params='IDFT')
    img_idft = np.dot(np.dot(Gl, img_dft), Gr) / (height * width)
    return img_idft


def dft_show(img_dft):
    """
    测试情况下显示图像用
    :param img_dft: DFT的变换结果，H*W，np.complex64
    :return:
    """
    dft_spectrum = np.abs(img_dft)  # 这里将complex64转换为float32
    dft_spectrum = (dft_spectrum / np.max(dft_spectrum) * 255)
    dft_spectrum = np.log(1 + dft_spectrum)
    dft_phase = np.angle(img_dft) + np.pi

    cv2.imshow('magn', np.clip(dft_spectrum, 0, 1))
    cv2.imshow('phase', dft_phase)

    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


def shift_spectrum_to_center(img):
    """
    将频谱原点移到图像中心 f(x,y)*(-1)^(x+y)
    :param img: 原始灰度图 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W
    :return: 变换后的,与原图大小类型均相同
    """
    img = np.float32(img)
    height, width = img.shape
    row = np.array([[math.pow(-1, x) for x in range(height)]])
    col = np.array([[math.pow(-1, y) for y in range(width)]])
    return img * (np.dot(row.T, col))


def dft_tool(img):
    """
    为可视化工具箱调用的DFT接口
    :param img: 仅限灰度图，形状H*W，数据类型numpy ndarray，范围0-255 uint8
    :return:
    """
    img = shift_spectrum_to_center(img.copy().astype(np.float32))
    img_dft = dft_complex(img)
    dft_spectrum = np.abs(img_dft)  # 这里将complex64转换为float32，显示频谱模长
    dft_spectrum = (dft_spectrum / np.max(dft_spectrum) * 255)  # 归一化
    dft_spectrum = np.log(1 + dft_spectrum)  # 这里做对数变化增强视觉效果
    np.clip(dft_spectrum, 0, 1)
    return (dft_spectrum * 255).astype(np.uint8), img_dft


def idft_tool(img_dft):
    """
    为可视化工具箱调用的IDFT接口
    :param img_dft: 形状H*W，数据类型numpy ndarray，np.complex64
    :return: 逆变换后的图，形状H*W，数据类型numpy ndarray，范围0-255 uint8
    """
    img_idft = idft_complex(img_dft)
    idft = np.abs(img_idft).astype(np.uint8)  # 只保留实部并取绝对值
    print(idft.shape)
    # idft = shift_spectrum_to_center(idft.copy()).astype(np.uint8)
    return idft


if __name__ == '__main__':
    src = cv2.imread('1.jpg')
    src = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    cv2.imshow('raw image', src)
    _, img_dft = dft_tool(src)
    dft_show(img_dft)
    img_idft = idft_tool(img_dft)
    cv2.imshow('idft', cv2.cvtColor(img_idft, cv2.COLOR_RGB2BGR))
    cv2.waitKey(0)
    cv2.destroyAllWindows()
